import pyaudio
import numpy as np
import time
import logging
import os
import pickle
import timeit
import warnings
import pyaudio
import requests

from main.feature_engineer import FeatureEngineer
from main.majority_voter import MajorityVoter
from main.baby_cry_predictor import BabyCryPredictor

# --- 模型地址 ---
LOAD_PATH_MODEL = os.path.normpath("./output/model")
SAVE_PATH = os.path.normpath("./output/prediction")
LOG_PATH = os.path.normpath("./logger")
        
# --- 实时音频处理设置 ---
CHUNK_DURATION_S = 2  # 每次回调的音频块时长（秒） - PyAudio 中对应 frames_per_buffer
BUFFER_DURATION_S = 4 # 内部处理缓冲区的总时长（秒）
SAMPLE_RATE = 44100 # 从 feature_extraction 导入
CHUNK_SAMPLES = int(CHUNK_DURATION_S * SAMPLE_RATE) # PyAudio 的 frames_per_buffer
BUFFER_SAMPLES = int(BUFFER_DURATION_S * SAMPLE_RATE)
FORMAT = pyaudio.paFloat32 # 数据格式，匹配 feature_extraction 中的 np.float32
CHANNELS = 1 # 单声道

# 使用一个环形缓冲区来存储最近的音频数据 (逻辑不变)
audio_buffer = np.zeros(BUFFER_SAMPLES, dtype=np.float32)
buffer_write_pos = 0

# 连续哭声阈值，用于去抖动
CONSECUTIVE_CRY_THRESHOLD = 2
cry_counter = 0
last_prediction_time = time.time()
# 预测频率每隔多少秒做一次预测
PREDICTION_INTERVAL_S = 5 

class RealmTimePrediction:
    def __init__(self):
        # 上次状态
        self.LAST_CRY_STATE = -1
        # 上次哭闹时间
        self.LAST_CRY_TIME = 0

    # --- PyAudio 回调函数 ---
    def pyaudio_callback(self, in_data, frame_count, time_info, status):
        """这是 PyAudio 在获取到新音频数据时调用的函数"""
        global audio_buffer, buffer_write_pos, cry_counter, last_prediction_time

        if status:
            print(f"PyAudio Status: {status}") # 打印潜在的音频流问题

        # 1. 将 PyAudio 传入的 bytes 数据转换为 numpy 数组
        # 确保 NumPy 数据类型与 PyAudio 格式匹配 (paFloat32 -> np.float32)
        try:
            audio_data = np.frombuffer(in_data, dtype=np.float32)
        except Exception as e:
            print(f"错误：转换音频数据时出错 - {e}")
            # 如果转换失败，不继续处理，并让流继续
            return (None, pyaudio.paContinue)

        # 2. 将新数据写入环形缓冲区 (逻辑不变)
        # frame_count 就是这次回调传入的样本数，应该等于 CHUNK_SAMPLES
        chunk_size = frame_count
        if buffer_write_pos + chunk_size < BUFFER_SAMPLES:
            audio_buffer[buffer_write_pos : buffer_write_pos + chunk_size] = audio_data
            buffer_write_pos += chunk_size
        else: # 处理环绕情况
            part1_size = BUFFER_SAMPLES - buffer_write_pos
            audio_buffer[buffer_write_pos:] = audio_data[:part1_size]
            part2_size = chunk_size - part1_size
            audio_buffer[:part2_size] = audio_data[part1_size:]
            buffer_write_pos = part2_size

        # 3. 控制预测频率 (逻辑不变)
        current_time = time.time()
        if current_time - last_prediction_time < PREDICTION_INTERVAL_S:
            # 时间未到，不做预测，但必须返回 paContinue 让流继续
            return (None, pyaudio.paContinue)

        last_prediction_time = current_time

        # 4. 准备用于预测的数据 (从环形缓冲区提取，逻辑不变)
        if buffer_write_pos == 0:
            analysis_data = audio_buffer[:]
        else:
            analysis_data = np.roll(audio_buffer, -buffer_write_pos)

        # 5. 提取特征 (逻辑不变)
        prediction = self.make_prediction(audio_data=analysis_data)
        
        # 6. 判断是否哭声
        majority_voter = MajorityVoter(prediction)
        majority_vote,lable = majority_voter.vote()
        
        # 7. 输出结果与去抖动 (逻辑不变)
        if majority_vote == 1: # 预测为哭声
            cry_counter += 1
            if cry_counter >= CONSECUTIVE_CRY_THRESHOLD:
                print(f"!!! 检测到婴儿哭声 !!! (PyAudio - 连续 {cry_counter} 次)")
                self.make_result(vote=1,lable=lable)
        else: # 预测为非哭声
            if cry_counter > 0:
                print(f"哭声停止或误报 (PyAudio - 之前连续计数: {cry_counter})")
                self.make_result(vote=0,lable=lable)
            cry_counter = 0
                
        # 9. PyAudio 回调函数必须返回一个元组 (out_data, flag)
        # 对于纯输入流，out_data 为 None
        # paContinue 表示继续处理音频流
        return (None, pyaudio.paContinue)

    # 最终输出结果
    def make_result(self, vote: int,lable: str):
        # 如果当前结果和上次不一致则状态发生改变，记录
        if (self.LAST_CRY_STATE != -1) and (self.LAST_CRY_STATE != vote):
            # 当前时间
            curr_cry_date = int(round(time.time()))
            cry_date = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(curr_cry_date))
            
            # 时长
            duration = curr_cry_date - self.LAST_CRY_TIME
            hour = duration // 3600
            min = (duration % 3600) // 60
            sec = duration % 60
            duration_str = f"{hour}时{min}分{sec}秒"
            
            # 预测结果
            pre_result = ""
            if vote == 1:
                descr = "好宝👍👍👍" if duration > 60 * 30 else "😒😒😒"
                pre_result = pre_result + f"{cry_date} 开始哭闹({lable})😭，安静时长:{duration_str} {descr}\n"
                # requests.get(url=f"https://api.day.app/oSiy4Mtp7WGTydYEvL7UtN/{pre_result}")
            else:
                descr = "😒😒😒" if duration > 60 * 30 else "好宝👍👍👍"
                pre_result = pre_result + f"{cry_date} 停止哭闹😊，哭闹时长:{duration_str}\n"
                # requests.get(url=f"https://api.day.app/oSiy4Mtp7WGTydYEvL7UtN/{pre_result}")
            
            print(pre_result)
            # 更新状态改变时间点
            self.LAST_CRY_TIME = int(round(time.time()))
         
            # 写入结果prediction.txt
            with open(os.path.join(SAVE_PATH, 'prediction.txt'), 'a', encoding='utf-8') as text_file:
                text_file.writelines(pre_result)
                print(pre_result)
        # 更新预测状态
        self.LAST_CRY_STATE= vote
        

    # 预测算法
    def make_prediction(self, audio_data):
        # 特征引擎
        logging.info('Starting feature engineering')
        start = timeit.default_timer()

        # 提取特征
        engineer = FeatureEngineer()
        play_list_processed = list()
        tmp = engineer.feature_engineer(audio_data)
        play_list_processed.append(tmp)

        stop = timeit.default_timer()
        logging.info('Time taken for feature engineering: {0}'.format(stop - start))

        # MAKE PREDICTION
        logging.info('Predicting...')
        start = timeit.default_timer()

        with warnings.catch_warnings():
            warnings.simplefilter("ignore", category=UserWarning)
            with open((os.path.join(LOAD_PATH_MODEL, 'model.pkl')), 'rb') as fp:
                model = pickle.load(fp)
        
        # 处理预测结果
        predictor = BabyCryPredictor(model)
        predictions = list()
        for signal, _ in play_list_processed:
            tmp = predictor.classify(signal)
            predictions.append(tmp)
        return predictions
        
    def real_time(self):
        # --- 启动音频流 ---
        p = pyaudio.PyAudio() # 初始化 PyAudio
        stream = None # 初始化 stream 变量

        print("使用 PyAudio 开始实时检测... 按 Ctrl+C 退出。")
        try:
            stream = p.open(format=FORMAT,
                            channels=CHANNELS,
                            rate=SAMPLE_RATE,
                            input=True, # 指定为输入流
                            frames_per_buffer=CHUNK_SAMPLES, # 每次回调处理的块大小
                            stream_callback=self.pyaudio_callback) # 设置回调函数

            stream.start_stream() # 开始音频流处理

            # 保持主线程运行，直到用户中断或流停止
            while stream.is_active():
                time.sleep(0.1)
        except KeyboardInterrupt:
            print("\n检测停止 (用户中断)。")
        except Exception as e:
            print(f"发生错误: {e}")
        finally:
            # --- 清理资源 ---
            if stream is not None:
                print("正在停止和关闭音频流...")
                if stream.is_active(): # 检查流是否仍在活动
                    stream.stop_stream()
                stream.close()
                print("音频流已关闭。")
            if p is not None:
                p.terminate() # 释放 PyAudio 资源
                print("PyAudio 实例已终止。")